#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/9/18 10:27
# @Author  : HYCX-AI-LAB Liang jinhao
import os

from utils.detection_3d.visualization import parse_npy
from utils.detection_3d.calculate_utils import *
import open3d as o3d
import numpy as np
import utils.detection_3d.calculate_side_volume as side_utils
from utils.MQ.message import send_message
import json


# import openpyxl


class Calculator():
    """3维点云数据体积计算

    """

    def __init__(self, MQ_send, connection_send, task_id, task_type, point_cloud_file, direction):
        """
        Args:
            MQ_send:
            connection_send:
            task_id:
            task_type:
            point_cloud_file: 文件路径
            direction: 数据方向，0是正面，1是侧面
        """
        # MQ
        self.MQ_send = MQ_send
        self.connection_send = connection_send

        self.task_id = task_id
        self.task_type = task_type

        self.point_cloud_file = point_cloud_file
        self.direction = direction

    def cal_front_volume(self):
        xyz_points = parse_npy(self.point_cloud_file)  # 得到x，y，z坐标
        x_points, y_points, z_points = xyz_points
        # o3d.visualization.draw_geometries([numpy2pcd(x_points, y_points, z_points)], window_name='AREA_1_TEP')

        # 未处理的数据显示

        bottom_down, bottom_up, bottom_left = get_bottom_value(x_points, y_points, z_points)
        Z_BOTTOM = (bottom_down + bottom_up + bottom_left) / 3  # PCD面的大致底面高度
        # ==========================AREA_1 最下面 第一个焊点的正面==============================
        # 根据区域，大致定位，过滤数据

        _x1, _y1, _z1 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 10 * 3 + 0.5,
                                            x_max=x_points.max() / 10 * 6 - 3.9, y_min=y_points.min() + 1,
                                            y_max=y_points.max() / 5 + 1)
        # o3d.visualization.draw_geometries([numpy2pcd(_x1, _y1, _z1)], window_name='AREA_1_TEP')
        # 根据底面高度，准确定位
        # draw_z_points(_x1, _z1)
        real_bottom = get_real_bottom(Z_BOTTOM, _z1)
        x1, y1, z1 = data_filter_by_Z_value(_x1, _y1, _z1, real_bottom)
        # draw_z_points(x1, z1)
        # o3d.visualization.draw_geometries([numpy2pcd(x1, y1, z1)], window_name='AREA_1')
        v1_front = np.sum(z1) * CONSTANT_UP
        print("the AREA_1 front volume is", v1_front)

        # ==========================AREA_2 第二个焊点的正面==============================
        # 根据区域，大致定位，过滤数据
        _x2, _y2, _z2 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 2 - 15,
                                            x_max=x_points.max() / 2 - 10, y_min=y_points.max() / 2 - 2,
                                            y_max=y_points.max() / 20 * 11)
        # o3d.visualization.draw_geometries([numpy2pcd(_x2, _y2, _z2)], window_name='AREA_2_TEP')
        # 根据底面高度，准确定位
        # draw_z_points(_x2, _z2)
        real_bottom = get_real_bottom(Z_BOTTOM, _z2)
        x2, y2, z2 = data_filter_by_Z_value(_x2, _y2, _z2, real_bottom)
        area_2_pcd = numpy2pcd(x2, y2, z2)
        # draw_z_points(x2, z2)
        # o3d.visualization.draw_geometries([area_2_pcd], window_name='AREA_2')
        _, _, z2 = pcd2numpy(area_2_pcd)
        v2_front = np.sum(z2) * CONSTANT_UP
        print("the AREA_2 front volume is", v2_front)

        # ==========================AREA_3 第三个焊点的正面==============================
        # 根据区域，大致定位，过滤数据
        _x3, _y3, _z3 = data_filter_by_area(x_points, y_points, z_points, x_min=x_points.max() / 8,
                                            x_max=x_points.max() / 8 * 3 - 3, y_min=y_points.max() / 5 * 4 - 1.5,
                                            y_max=y_points.max() - 1.5)
        # o3d.visualization.draw_geometries([numpy2pcd(_x3, _y3, _z3)], window_name='AREA_3_TEP')
        # 根据底面高度，准确定位
        # draw_z_points(_x3, _z3)
        real_bottom = get_real_bottom(Z_BOTTOM, _z3)
        x3, y3, z3 = data_filter_by_Z_value(_x3, _y3, _z3, real_bottom)
        # draw_z_points(x3, z3)
        area_3_pcd = numpy2pcd(x3, y3, z3)

        # o3d.visualization.draw_geometries([area_3_pcd], window_name='AREA_3')
        _, _, z3_front_res = pcd2numpy(area_3_pcd)
        v3_front = np.sum(z3_front_res) * CONSTANT_UP
        print("the AREA_3 front volume is", v3_front)

        return v1_front, v2_front, v3_front

    # def cal_side_volume(self):
    #     xyz_points = parse_npy(self.point_cloud_file)
    #     x_points, y_points, z_points = xyz_points
    #
    #     # 得到包括基面和pcd面的数据
    #     _x, _y, _z = data_filter_by_area(x_points, y_points, z_points, y_min=5, y_max=y_points.max() - 3.5)  # 去除y最下面的干扰
    #     _x, _y, _z = denoising_by_remove_points(_x, _y, _z, show=False)
    #     # draw_z_points(_y,_z)
    #     # 三次去基面逼近PIN
    #     base1 = get_real_bottom(_z.max(), _z) - 2.5
    #     _x, _y, _z = data_filter_by_Z_value(_x, _y, _z, Z_BOTTOM=base1, compare=-1)
    #
    #     base2 = get_real_bottom(_z.max(), _z) - 2.5
    #     _x, _y, _z = data_filter_by_Z_value(_x, _y, _z, Z_BOTTOM=base2, compare=-1)
    #
    #     base3 = get_real_bottom(_z.max(), _z) - 1.5
    #     _x, _y, _z = data_filter_by_Z_value(_x, _y, _z, Z_BOTTOM=base3, compare=-1)
    #
    #     # o3d.visualization.draw_geometries([numpy2pcd(_x, _y, _z)], window_name='Interested_AREA')
    #
    #     # 4.5为下基面
    #     _x, _y, _z = data_filter_by_Z_value(_x, _y, _z, Z_BOTTOM=-4.0, compare=1)
    #     _x, _y, _z = denoising_by_remove_points(_x, _y, _z, show=False)
    #     # o3d.visualization.draw_geometries([numpy2pcd(_x, _y, _z)], window_name='Interested_AREA')
    #
    #     # ==========================AREA_1 ==============================
    #     # 第一个焊面,三维图最下面的锡焊：包括两个焊平面，分明是PCD正面的侧平面以及夹层的侧平面
    #
    #     x1, y1, z1 = data_filter_by_area(_x, _y, _z, y_max=_y.min() + 4)
    #
    #     # 去除一些离散点
    #     x1, y1, z1 = denoising_by_remove_points(x1, y1, z1, show=False)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x1, y1, z1)], window_name='AREA_1')
    #
    #     # 计算PIN针的深度
    #     """
    #     根据最高点确定
    #     """
    #     PIN_1_CENTER_POINT = get_PIN_line_value(x1, y1, z1)
    #     # 获取侧面点云
    #     x1, y1, z1 = data_filter_by_area(x1, y1, z1, x_max=x1.min() + 3.3)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x1, y1, z1)], window_name='AREA_1_SIDE_WELDING')
    #     AREA_1_PIN_Z_VALUE = PIN_1_CENTER_POINT[0]
    #     z1 = z1[z1 > AREA_1_PIN_Z_VALUE]
    #     v1_side = np.sum(z1 - AREA_1_PIN_Z_VALUE) * CONSTANT
    #     print("the AREA_1 side volume is", v1_side)
    #
    #     # # 绘画底面查看是否准确 z1 = z1[z1>AREA_1_PIN_Z_VALUE]要注释掉
    #     # x1, y1, z1 = add_bottom(x1, y1, z1, AREA_1_PIN_Z_VALUE)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x1,y1,z1)], window_name='REAL_AREA_1')
    #
    #     # ==========================AREA_2 ==============================
    #
    #     Y2_MIN = _y.min() + 9
    #     Y2_MAX = _y.min() + 16
    #
    #     x2, y2, z2 = data_filter_by_area(_x, _y, _z, y_min=Y2_MIN, y_max=Y2_MAX)
    #     # # 去除一些离散点
    #     x2, y2, z2 = denoising_by_remove_points(x2, y2, z2, show=False)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x2, y2, z2)], window_name='AREA_2')
    #
    #     # 计算PIN针的深度
    #     PIN_2_CENTER_POINT = get_PIN_line_value(x2, y2, z2)
    #     # 获取侧面点云
    #     x2, y2, z2 = data_filter_by_area(x2, y2, z2, x_max=x2.min() + 3.3)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x2, y2, z2)], window_name='AREA_2_SIDE_WELDING')
    #     AREA_2_PIN_Z_VALUE = PIN_2_CENTER_POINT[0]
    #
    #     z2 = z2[z2 > AREA_2_PIN_Z_VALUE]
    #     v2_side = np.sum(z2 - AREA_2_PIN_Z_VALUE) * CONSTANT
    #     print("the AREA_2 side volume is", v2_side)
    #     # # 绘画底面查看是否准确
    #     # x2, y2, z2 = add_bottom(x1, y1, z1, AREA_2_PIN_Z_VALUE)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x2, y2, z2)], window_name='REAL_AREA_2')
    #
    #     # ==========================AREA_3 ==============================
    #
    #     Y3_MIN = _y.min()
    #     Y3_MAX = _y.min() + 6
    #
    #     x3, y3, z3 = data_filter_by_area(_x, _y, _z, y_min=Y3_MIN, y_max=Y3_MAX)
    #
    #     # 去除一些离散点
    #     x3, y3, z3 = denoising_by_remove_points(x3, y3, z3, show=False)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x3, y3, z3)], window_name='AREA_3')
    #     # 计算PIN针的深度
    #
    #     PIN_3_CENTER_POINT = get_PIN_line_value(x3, y3, z3)
    #     # 获取侧面点云
    #     x3, y3, z3 = data_filter_by_area(x3, y3, z3, x_max=x3.min() + 3.3)
    #     # o3d.visualization.draw_geometries([numpy2pcd(x3, y3, z3)], window_name='AREA_3_SIDE_WELDING')
    #     AREA_3_PIN_Z_VALUE = PIN_3_CENTER_POINT[0]
    #     # print(AREA_3_PIN_Z_VALUE)
    #     z3 = z3[z3 > AREA_3_PIN_Z_VALUE]
    #     v3_side = np.sum(z3 - AREA_3_PIN_Z_VALUE) * CONSTANT
    #     print("the AREA_3 side volume is", v3_side)
    #     return v1_side, v2_side, v3_side

    def cal_side_volume(self):
        # 清理nan数据
        print(self.point_cloud_file)
        npy_data = side_utils.clean_height(self.point_cloud_file)
        # side_utils.vis_pcd(side_utils.read_3d_npy(npy_data)[0])
        npy_data = side_utils.filter_pin_npy(npy_data, 5, 2)
        # side_utils.vis_pcd(side_utils.read_3d_npy(npy_data)[0])
        # 提取pin针
        area = side_utils.extract_valid_area(npy_data)
        pin1 = npy_data[area[0][0]:area[0][1], :]
        pin2 = npy_data[area[1][0]:area[1][1], :]
        pin3 = npy_data[area[2][0]:area[2][1], :]
        # side_utils.vis_pcd(side_utils.read_3d_npy(pin1)[0])
        # side_utils.vis_pcd(side_utils.read_3d_npy(pin2)[0])
        # side_utils.vis_pcd(side_utils.read_3d_npy(pin3)[0])
        # 清理离散点
        pin1 = side_utils.clean_discrete_points(pin1, 2)
        pin2 = side_utils.clean_discrete_points(pin2, 2)
        pin3 = side_utils.clean_discrete_points(pin3, 2)
        # pin针上下分离
        pin1_down, pin1_up = side_utils.divide_up_side_pin(pin1)
        pin2_down, pin2_up = side_utils.divide_up_side_pin(pin2)
        pin3_down, pin3_up = side_utils.divide_up_side_pin(pin3)
        # pin针反射值清理
        index = side_utils.clean_reflection(pin1_down)
        pin1 = pin1[index:, :]
        pin1_down = pin1_down[index:, :]
        pin1_up = pin1_up[index:, :]
        index = side_utils.clean_reflection(pin2_down)
        pin2 = pin2[index:, :]
        pin2_down = pin2_down[index:, :]
        pin2_up = pin2_up[index:, :]
        index = side_utils.clean_reflection(pin3_down)
        pin3 = pin3[index:, :]
        pin3_down = pin3_down[index:, :]
        pin3_up = pin3_up[index:, :]

        # 粗略估计pin针深度
        depth_pin1 = side_utils.cut_slice(pin1, pin1_up, pin1_down, 1)
        depth_pin2 = side_utils.cut_slice(pin2, pin2_up, pin2_down, 2)
        depth_pin3 = side_utils.cut_slice(pin3, pin3_up, pin3_down, 3)

        # 获得精确深度
        real_pin1 = side_utils.get_real_side_depth(pin1, depth_pin1)
        real_pin2 = side_utils.get_real_side_depth(pin1, depth_pin2)
        real_pin3 = side_utils.get_real_side_depth(pin1, depth_pin3)
        depth_pin1 = min(depth_pin1 + 0.2, max(depth_pin1, real_pin1))
        depth_pin2 = min(depth_pin2 + 0.2, max(depth_pin2, real_pin2))
        depth_pin3 = min(depth_pin2 + 0.2, max(depth_pin3, real_pin3))

        # 计算pin针体积
        pin1_val = side_utils.calculate_side_volumn_pin(pin1_down, depth_pin1)
        # vis_data, x, y, z = side_utils.read_3d_npy_add_background(pin1, depth_pin1)
        # side_utils.vis_pcd(vis_data)
        pin2_val = side_utils.calculate_side_volumn_pin(pin2_down, depth_pin2)
        # vis_data, x, y, z = side_utils.read_3d_npy_add_background(pin2, depth_pin2)
        # side_utils.vis_pcd(vis_data)
        pin3_val = side_utils.calculate_side_volumn_pin(pin3_down, depth_pin3)
        # vis_data, x, y, z = side_utils.read_3d_npy_add_background(pin3, depth_pin3)
        # side_utils.vis_pcd(vis_data)
        return pin1_val, pin2_val, pin3_val

    # new\pin3_35051772633505277134A02092200190226_side_2022-09-02-06-15-25.png
    # new\pin3_35051772633505277134A02092200190227_side_2022-09-02-06-16-08.png
    def send_res(self):
        Pin_1_V, Pin_2_V, Pin_3_V = 0.0, 0.0, 0.0
        try:
            if self.direction == 0:
                Pin_1_V, Pin_2_V, Pin_3_V = self.cal_front_volume()
            else:
                Pin_1_V, Pin_2_V, Pin_3_V = self.cal_side_volume()
        except ValueError as e:
            print("the data may be none")
        finally:
            res = {
                "task_id": self.task_id,
                "task_type": "3d_volume",
                "point_cloud_file": self.point_cloud_file,
                "direction": self.direction,
                "Pin_1_V": Pin_1_V,
                "Pin_2_V": Pin_2_V,
                "Pin_3_V": Pin_3_V,
            }
            send_message(self.MQ_send, self.connection_send, json.dumps(res, ensure_ascii=False))

    def write_result(self, target, if_write=False):
        import openpyxl
        if self.point_cloud_file.endswith('.npy') and 'up' in self.point_cloud_file:
            Pin_1_V, Pin_2_V, Pin_3_V = self.cal_front_volume()
        else:
            Pin_1_V, Pin_2_V, Pin_3_V = self.cal_side_volume()
        if if_write:
            if os.path.exists(target):
                data = openpyxl.load_workbook(target)
            else:
                data = openpyxl.Workbook()
            stn = data.sheetnames[0]
            table = data.active
            nrows = table.max_row
            table.cell(nrows + 1, 1).value = self.point_cloud_file
            table.cell(nrows + 1, 2).value = Pin_1_V
            table.cell(nrows + 1, 3).value = Pin_2_V
            table.cell(nrows + 1, 4).value = Pin_3_V
            data.save(target)


if __name__ == "__main__":
    """
      计算正面
    """
    # cal = Calculator(0, 0, 0, 0,
    #                  "D:\\Programs\\data\\dataset\\20220831\\35051772633505277134A31082200188915_up_2022-08-31-09-45-59.npy",
    #                  0)
    # cal.send_res()
    # cal = Calculator(0, 0, 0, 0,
    #                  "D:\\Programs\\data\\dataset\\20220831\\35051772633505277134A31082200188869_side_2022-08-31-09-00-10.npy",
    #                  1)
    # cal.send_res()
    target = "D:\\Programs\\data\\result.xlsx"
    dataRoot = "D:\\Programs\\data\\dataset\\20220902"
    file = os.listdir(dataRoot)
    for onefile in file:
        if (onefile[-3:] == 'npy' and onefile.count("_side_") > 0):
            cal = Calculator(0, 0, 0, 0, os.path.join(dataRoot, onefile), 0)
            # cal.write_result(target, True)
